# Arman and Seamus performance harness...

import os

class CgroupAdmin(object):
	
	def __init__(self):
		self.cgroup_name = ""
		self.process_name = ""
		self.memory_in_bytes = 0
		self.cpu_shares = 0

	def get_cgroup_values(self):
		print "--- CGroup Admin Program ---"
		print "This program will create a cgroup for a specific process and set the memory usage."
		try:
			local_cgroup_name = raw_input('Please enter the name for this cgroup: ')
			local_process_name = raw_input('Please enter the name of the process to add to this cgroup: ')
			local_memory_in_bytes = int(raw_input('Please enter value in for the memory limit for this process: '))
			local_cpu_shares = int(raw_input('Please enter value for the cpu shares for this process: '))
			if local_cgroup_name == "":
				print "The cgroup name cannot be null!"
				exit(1)
			if local_process_name == "":
				print "The cgroup name cannot be null!"
				exit(1)	
			if local_memory_in_bytes == 0:
				print "The memory usage cannot be 0!"
				exit(1)		
			if local_cpu_shares == 0:
				print "The cpu shares cannot be 0!"
				exit(1)		
			self.cgroup_name = local_cgroup_name 
			self.process_name = local_process_name 
			self.memory_in_bytes = local_memory_in_bytes
			self.cpu_shares = local_cpu_shares
		except Exception as e:
			"Error getting cgroup values..." + e.message
		
	def create_cgroup(self):
		try:
			os.system('cgcreate -g memory:' + self.cgroup_name)
			os.system('cgcreate -g cpu:' + self.cgroup_name)
			os.system('cgclassify --sticky -g memory:' + self.cgroup_name + ' $(pidof ' + self.process_name  + ')')
			os.system('cgclassify --sticky -g cpu:' + self.cgroup_name + ' $(pidof ' + self.process_name  + ')')
			os.system('cd /sys/fs/cgroup/memory/' + self.cgroup_name)
			os.system(' echo ' + str(self.memory_in_bytes) + ' > memory.limit_in_bytes')
			os.system('cd /sys/fs/cgroup/cpu/' + self.cgroup_name)
			os.system(' echo ' + str(self.cpu_shares) + ' > cpu.shares')
			print 'cgroup ' + self.cgroup_name + ' created successfully!'
		except Exception as e:
			"Error creating cgroup..." + e.message

def main():
    try:
    	if not os.geteuid() == 0: 
    		print "This script must run as root!"
    		exit(1)
        ca = CgroupAdmin()
        ca.get_cgroup_values()
        ca.create_cgroup()
        print "Processing complete."
    except Exception as e:
        print "Error " + e.message
        exit(1)


if __name__ == "__main__":
    main()